///Program stitching images together

//Standard libraries
#include <string>
#include <fstream>
#include <vector>
#include <sys/types.h>
#include <sys/stat.h>

//Boost libraries
#include <boost/operators.hpp>
#include <boost/program_options.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/convenience.hpp>
#include <boost/filesystem/fstream.hpp>

//VisionWorkbench libraries
#include <vw/InterestPoint.h>
#include <vw/Math.h>
#include <vw/Image/ImageView.h>
#include <vw/Image/ImageViewRef.h>
#include <vw/Image/BlockRasterize.h>
#include <vw/Core/Debugging.h>
#include <vw/FileIO.h>
#include <vw/FileIO/DiskImageView.h>
#include <vw/Mosaic/QuadTreeGenerator.h>
#include <vw/Mosaic/ImageComposite.h>

#define MAX_POINTS_TO_DRAW 1000

using namespace vw;
using namespace vw::ip;
namespace po = boost::program_options;
namespace po = boost::program_options;
namespace fs = boost::filesystem;

typedef std::vector<std::pair<int, Matrix<double> > > BestMatchingTransforms;
typedef std::pair<std::vector<InterestPoint>, std::vector<InterestPoint> > MatchPoints;
typedef DiskImageView<PixelRGBA<float> > UsedImageType;

class ImageDataCollector{
    private:
        std::string mosaic_name; 
        std::string interest_operator;
        std::string descriptor_generator;
        float harris_threshold, log_threshold;
        int tile_size;
        bool single_scale;
        bool verbose;
        bool debug;

        static std::string prefix_from_filename(std::string const& filename) {
            std::string result = filename;
            int index = result.rfind(".");
            if (index != -1) 
                result.erase(index, result.size());
            return result;
        }

        /// Erases a file suffix if one exists and returns the base string
        static std::string suffix_from_filename(std::string const& filename) {
            std::string result = filename;
            int index = result.rfind(".");
            if (index != -1) 
                result.erase(0, index);
            return result;
        }

        // Draw the interest points and write as an image.
        template <class ViewT>
            void write_point_image(std::string out_file_name, 
                    ImageViewBase<ViewT> const& src,
                    InterestPointList const& points) {

                ImageView<PixelRGB<uint8> > viz = pixel_cast<PixelRGB<uint8> >(channel_cast_rescale<uint8>(src));

                // Draw points into color planes
                int n = 0;
                for (InterestPointList::const_iterator pt = points.begin();
                        pt != points.end() && n < MAX_POINTS_TO_DRAW; ++pt, ++n) {
                    // Draw a red line from the point outward along the orientation
                    for (int r=0; r<(int)(8*(*pt).scale); ++r){
                        int i = (int)(0.5 + (*pt).x + r*cos((*pt).orientation));
                        int j = (int)(0.5 + (*pt).y + r*sin((*pt).orientation));
                        // red, green, blue
                        viz(i,j) = PixelRGB<uint8>(255, 0, 0);
                    }
                    // Draw a green 3x3 filled square at the point to indicate center
                    int i0 = (int)(0.5 + (*pt).x);
                    int j0 = (int)(0.5 + (*pt).y);
                    for (int j=j0-1; j<=j0+1; ++j){
                        for (int i=i0-1; i<=i0+1; ++i){
                            // red, green, blue
                            viz(i,j) = PixelRGB<uint8>(0, 255, 0);
                        }
                    }
                }

                write_image(out_file_name, viz);
            }

    public:
        //Freely accessible containers
        std::map<std::string, fs::path> image_file_names;           //Image filenames on the disk
        std::map<std::string, UsedImageType > image_files;          //Images
        std::map<std::string, InterestPointList> interest_points;   //Interest points
        std::map<std::string, BestMatchingTransforms > bestMatches; //Best matches of an image


        ImageDataCollector(std::string mosaic_name, bool debug){
            this->mosaic_name = mosaic_name;
            //Default values:
            tile_size = 512;
            interest_operator = "LoG";
            single_scale = false;
            log_threshold = 0.03;
            harris_threshold = 1e-5;
            this->debug = debug;
            descriptor_generator = "patch";
        }

        void extractFeatures(){
            // Collect paths to images in the directory
            fs::path source_dir_path( mosaic_name, fs::native );
            fs::directory_iterator pi( source_dir_path ), pend;
            for( ; pi != pend; ++pi ) {
                image_file_names[basename(*pi)] = *pi;
            }

            // Load the images
            std::map<std::string,fs::path>::iterator ifi=image_file_names.begin(), ifend=image_file_names.end();

            //Pointers to last inserted image and interest points to the maps, in order not to search by name every time
            std::pair<   std::map<  std::string, UsedImageType  >::iterator, bool   > lastInsertedImage;
            std::pair<   std::map<  std::string, InterestPointList >::iterator, bool   > lastInsertedPoints;

            for( ; ifi != ifend; ++ifi ) {
                vw_out(InfoMessage)<< "Importing image file " << ifi->second.string() << std::endl;
                lastInsertedImage = image_files.insert(std::pair<std::string, UsedImageType >(ifi->first, UsedImageType(ifi->second.string())));

                vw_out(InfoMessage)<< "Detecting its interest points" << std::endl;

                if (interest_operator == "Harris") {
                    HarrisInterestOperator interest_operator(harris_threshold);
                    if (!single_scale) {
                        ScaledInterestPointDetector<HarrisInterestOperator> detector(interest_operator);
                        lastInsertedPoints = interest_points.insert(std::pair<std::string, InterestPointList>(ifi->first, detect_interest_points(lastInsertedImage.first->second, detector)));
                    } else {
                        InterestPointDetector<HarrisInterestOperator> detector(interest_operator);
                        lastInsertedPoints = interest_points.insert(std::pair<std::string, InterestPointList>(ifi->first, detect_interest_points(lastInsertedImage.first->second, detector)));
                    }
                } else if (interest_operator == "LoG") {
                    // Use a scale-space Laplacian of Gaussian feature detector. The
                    // associated threshold is abs(interest) > interest_threshold.
                    LogInterestOperator interest_operator(log_threshold);
                    if (!single_scale) {
                        ScaledInterestPointDetector<LogInterestOperator> detector(interest_operator);
                        lastInsertedPoints = interest_points.insert(std::pair<std::string, InterestPointList>(ifi->first, detect_interest_points(lastInsertedImage.first->second, detector)));
                    } else {
                        InterestPointDetector<LogInterestOperator> detector(interest_operator);

                        lastInsertedPoints = interest_points.insert(std::pair<std::string, InterestPointList>(ifi->first, detect_interest_points(lastInsertedImage.first->second, detector)));
                    }
                }

                vw_out(InfoMessage)<< "Found " << lastInsertedPoints.first->second.size()<<" points"<<std::endl;

                // Write out images with interest points marked
                if (debug) {
                    std::string prefix = prefix_from_filename(ifi->second.string());
                    std::string suffix = suffix_from_filename(ifi->second.string());
                    vw_out(InfoMessage)<<"Writing marked interest points to "<< prefix+"-IP" +suffix<<std::endl;
                    write_point_image(prefix + "-IP" + suffix, lastInsertedImage.first->second, lastInsertedPoints.first->second);
                }

                // Generate descriptors for interest points.
                vw_out(InfoMessage) << "Generating descriptors using " << descriptor_generator << " generator... ";
                if (descriptor_generator == "patch") {
                    PatchDescriptorGenerator descriptor;
                    descriptor(lastInsertedImage.first->second, lastInsertedPoints.first->second);
                }
                else {
                    //TODO: sift,what are the exr files?
                    PCASIFTDescriptorGenerator descriptor("pca_basis.exr", "pca_avg.exr");
                    descriptor(lastInsertedImage.first->second, lastInsertedPoints.first->second);
                }

                vw_out(InfoMessage) << "done.\n";

            }

        }
};


class ImageMatcher{

    private:
        static std::string prefix_from_filename(std::string const& filename) {
            std::string result = filename;
            int index = result.rfind(".");
            if (index != -1) 
                result.erase(index, result.size());
            return result;
        }

        /// Erases a file suffix if one exists and returns the base string
        static std::string suffix_from_filename(std::string const& filename) {
            std::string result = filename;
            int index = result.rfind(".");
            if (index != -1) 
                result.erase(0, index);
            return result;
        }

        // Draw the two images side by side with matching interest points
        // shown with lines.
        template <class ViewT>
            void write_match_image(std::string out_file_name, 
                    ImageViewBase<ViewT> const& src1,
                    ImageViewBase<ViewT> const& src2,
                    std::vector<InterestPoint> matched_ip1, 
                    std::vector<InterestPoint> matched_ip2) {
                mosaic::ImageComposite<PixelRGB<uint8> > composite;
                composite.insert(pixel_cast<PixelRGB<uint8> >(channel_cast_rescale<uint8>(src1.impl())),0,0);
                composite.insert(pixel_cast<PixelRGB<uint8> >(channel_cast_rescale<uint8>(src2.impl())),src1.impl().cols(),0);
                composite.set_draft_mode( true );
                composite.prepare();

                // Rasterize the composite so that we can draw on it.
                ImageView<PixelRGB<uint8> > comp = composite;

                // Draw a red line between matching interest points
                for (unsigned int i = 0; i < matched_ip1.size(); ++i) {
                    Vector2 start(matched_ip1[i].x, matched_ip1[i].y);
                    Vector2 end(matched_ip2[i].x+src1.impl().cols(), matched_ip2[i].y);
                    for (float r=0; r<1.0; r+=1/norm_2(end-start)){
                        int i = (int)(0.5 + start.x() + r*(end.x()-start.x()));
                        int j = (int)(0.5 + start.y() + r*(end.y()-start.y()));
                        // red, green, blue
                        comp(i,j) = PixelRGB<uint8>(255, 0, 0);
                    }
                }

                write_image(out_file_name, comp);
            }

        // --------------------------------------------------------------------------------------

        float matcher_threshold;
        ImageDataCollector *collector;
        bool verbose,debug, homography;
        int n_best_matches;

    public:
        ImageMatcher(ImageDataCollector &collector, bool debug){
            this->collector = &collector;
            // Set Vision Workbench debug output level
            if ( verbose != 0 )
                set_debug_level(DebugMessage);
            else 
                set_debug_level(InfoMessage);

            this->debug = debug;
            homography = true;

            matcher_threshold = 0.5;

            n_best_matches = 4;
        }

        MatchPoints align(std::string leftImage, std::string rightImage) {

            std::cout<< "Aligning " << leftImage << " and " << rightImage << std::endl;
            std::map<  std::string, UsedImageType >::iterator left_image, right_image;
            std::map<  std::string, InterestPointList >::iterator left_image_ip, right_image_ip;

            // Find the two images
            left_image = collector->image_files.find(leftImage);
            right_image = collector->image_files.find(rightImage);

            if (left_image == collector->image_files.end() || right_image == collector->image_files.end()){
                throw std::string("Error: images not loaded");
            }

            // Load their interest points
            left_image_ip = collector->interest_points.find(leftImage);
            right_image_ip = collector->interest_points.find(rightImage);

            if (left_image_ip == collector->interest_points.end() || right_image_ip == collector->interest_points.end()){
                throw std::string("Error: interest points not loaded");
            }

            // Image Alignment
            //
            // Images are aligned by computing interest points, matching
            // them using a standard 2-Norm nearest-neighor metric, and then
            // rejecting outliers by fitting a similarity between the
            // putative matches using RANSAC.

            // The basic interest point matcher does not impose any
            // constraints on the matched interest points.
            vw_out(InfoMessage) << "\nInterest Point Matching:\n";

            // RANSAC needs the matches as a vector, and so does the matcher.
            // this is messy, but for now we simply make a copy.
            std::vector<InterestPoint> ip1_copy(left_image_ip->second.size()), ip2_copy(right_image_ip->second.size());
            std::copy(left_image_ip->second.begin(), left_image_ip->second.end(), ip1_copy.begin());
            std::copy(right_image_ip->second.begin(), right_image_ip->second.end(), ip2_copy.begin());

            DefaultMatcher matcher(matcher_threshold);
            std::vector<InterestPoint> matched_ip1, matched_ip2;
            matcher(ip1_copy, ip2_copy, matched_ip1, matched_ip2, false, TerminalProgressCallback());
            vw_out(InfoMessage) << "\tFound " << matched_ip1.size() << " putative matches.\n";

            return MatchPoints(matched_ip1, matched_ip2);
        }


        void align_all(){
            std::cout<< "Aligning all image pairs"<<std::endl;
            std::map<std::string,fs::path>::iterator leftImage;
            for (leftImage = collector->image_file_names.begin(); leftImage != collector->image_file_names.end(); ++leftImage){
                std::map<std::string,fs::path>::iterator rightImage = leftImage;
                rightImage = leftImage;

                std::pair<std::map<std::string, BestMatchingTransforms>::iterator, bool> bmti = 
                    collector->bestMatches.insert(std::pair<std::string, BestMatchingTransforms>(leftImage->first, std::vector<std::pair<int, Matrix<double> > >()));
                for (++rightImage; rightImage != collector->image_file_names.end(); ++rightImage){
                    MatchPoints matchPoints = align(leftImage->first, rightImage->first);

                    // Write out the putative point correspondence image
                    std::string prefix = prefix_from_filename(leftImage->second.string())+"-"+prefix_from_filename(rightImage->first);
                    std::string suffix = suffix_from_filename(leftImage->second.string());

                    //  output putative match
                    std::map<  std::string, UsedImageType >::iterator left_image, right_image;
                    // Find the two images
                    left_image = collector->image_files.find(leftImage->first);
                    right_image = collector->image_files.find(rightImage->first);
                    if (debug){
                        write_match_image(prefix+"-putative-match" + suffix, left_image->second, right_image->second, matchPoints.first, matchPoints.second);
                    }

                    // RANSAC is used to fit a similarity transform between the
                    // matched sets of points
                    try{
                        std::vector<Vector3> ransac_ip1 = vw::ip::iplist_to_vectorlist(matchPoints.first);
                        std::vector<Vector3> ransac_ip2 = vw::ip::iplist_to_vectorlist(matchPoints.second);
                        Matrix<double> align_matrix;
                        if (homography) 
                            align_matrix = vw::math::ransac(ransac_ip2, ransac_ip1, 
                                    vw::math::HomographyFittingFunctor(),
                                    vw::math::InterestPointErrorMetric());
                        else
                            align_matrix = vw::math::ransac(ransac_ip2, ransac_ip1, 
                                    vw::math::AffineFittingFunctor(),
                                    vw::math::InterestPointErrorMetric());

                        std::cout << align_matrix << std::endl;

                        vw_out(InfoMessage) << "Storing the transoformation matrix for " << left_image->first << " - " << right_image->first << std::endl;
                        bmti.first->second.push_back(std::pair<int, Matrix<double> >(matchPoints.first.size(), align_matrix));

                        /*
                           vw_out(InfoMessage) << "Writing out aligned pair of images\n";

                        // Write out the aligned pair of images
                        ImageViewRef<PixelRGB<uint8> > aligned_image = transform(right_image->second, HomographyTransform(align_matrix),
                        left_image->second.cols(), left_image->second.rows());
                        write_image(prefix+".jpg", aligned_image); 
                         */
                    }
                    catch (vw::math::RANSACErr e){
                        std::cout << "No match between images\n" << std::endl;
                    }
                    catch (std::string errorMessage){
                        std::cout << errorMessage << std::endl;
                    }



                }
            }

        }
};

class ImageComposer{

    //parameters
    private:
        std::string mosaic_name;
        std::string file_type;
        int tile_size;
        bool draft;
        bool qtree;
        bool grayscale;

    public:
        ImageComposer(std::string mosaic_name, std::string file_type,
                unsigned int cache_size, int tile_size, bool draft,
                bool qtree, bool verbose, bool grayscale){

            if (verbose)
                set_debug_level(VerboseDebugMessage);
            vw_system_cache().resize( cache_size*1024*1024 );

            this->tile_size = tile_size;
            this->qtree = qtree;
            this->grayscale = grayscale;
            this->draft = draft;
        }

        void blend() {
            if( grayscale ) {
                do_blend<PixelGrayA<float> >();
            }
            else {
                do_blend<PixelRGBA<float> >();
            }
        }

        template <class PixelT>
            void do_blend() {
                mosaic::ImageComposite<PixelT> composite;
                if( draft ) composite.set_draft_mode( true );

                std::map<std::string,fs::path> image_files;
                std::map<std::string,fs::path> offset_files;
                fs::path source_dir_path( mosaic_name, fs::native );
                fs::directory_iterator pi( source_dir_path ), pend;
                for( ; pi != pend; ++pi ) {
                    if( extension(*pi) == ".offset" )
                        offset_files[basename(*pi)] = *pi;
                    else image_files[basename(*pi)] = *pi;
                }
                std::map<std::string,fs::path>::iterator ofi=offset_files.begin(), ofend=offset_files.end();
                for( ; ofi != ofend; ++ofi ) {
                    std::map<std::string,fs::path>::iterator ifi = image_files.find( ofi->first );
                    if( ifi != image_files.end() ) {
                        fs::ifstream offset( ofi->second );
                        int x, y;
                        offset >> x >> y;
                        std::cout << "Importing image file " << ifi->second.string() 
                            << " at offet (" << x << "," << y << ")" << std::endl;
                        composite.insert( DiskImageView<PixelT>( ifi->second.string(), false ), x, y );
                    }
                }

                vw_out(InfoMessage) << "Preparing the composite..." << std::endl;
                composite.prepare();
                if( qtree ) {
                    vw_out(InfoMessage) << "Preparing the quadtree..." << std::endl;
                    mosaic::QuadTreeGenerator quadtree( composite, mosaic_name );
                    quadtree.set_file_type( file_type );
                    quadtree.set_tile_size( tile_size );
                    vw_out(InfoMessage) << "Generating..." << std::endl;
                    quadtree.generate();
                    vw_out(InfoMessage) << "Done!" << std::endl;
                }
                else {
                    vw_out(InfoMessage) << "Blending..." << std::endl;
                    write_image( mosaic_name+".blend."+file_type, composite );
                    vw_out(InfoMessage) << "Done!" << std::endl;
                }
            }

};

int main( int argc, char *argv[] ) {
    std::string mosaic_name;
    std::string file_type;
    unsigned int cache_size;
    int tile_size;
    bool draft;
    bool qtree;
    bool verbose;
    bool grayscale;
    bool debug;

    po::options_description desc("Options");
    desc.add_options()
        ("help", "Display this help message")
        ("input-dir", po::value<std::string>(&mosaic_name), 
         "Explicitly specify the input directory")
        ("file-type", po::value<std::string>(&file_type)->default_value("png"), 
         "Output file type")
        ("tile-size", po::value<int>(&tile_size)->default_value(256), 
         "Tile size, in pixels")
        ("cache", po::value<unsigned>(&cache_size)->default_value(1024), 
         "Cache size, in megabytes")
        ("draft", "Draft mode (no blending)")
        ("qtree", "Output in quadtree format")
        ("grayscale", "Process in grayscale only")
        ("debug", "Output debug images")
        ("verbose", "Verbose output");
    po::positional_options_description p;
    p.add("input-dir", 1);

    po::variables_map vm;
    po::store( po::command_line_parser( argc, argv ).options(desc).positional(p).run(), vm );
    po::notify( vm );

    if( vm.count("help") ) {
        std::cout << desc << std::endl;
        return 1;
    }

    if( vm.count("input-dir") != 1 ) {
        std::cout << "Error: Must specify one (and only one) input directory!" << std::endl;
        std::cout << desc << std::endl;
        return 1;
    }

    if( vm.count("verbose") ) verbose = true; else verbose = false;
    if( vm.count("draft") ) draft = true; else draft = false;
    if( vm.count("qtree") ) qtree = true; else qtree = false;
    if( vm.count("grayscale") ) grayscale = true; else grayscale = false;
    if( vm.count("debug") ) debug = true; else debug = false;

    if( tile_size <= 0 ) {
        std::cerr << "Error: The tile size must be a positive number!  (You specified " 
            << tile_size << ".)" << std::endl;
        std::cout << desc << std::endl;
        return 1;
    }

    ImageDataCollector collector(mosaic_name, debug);
    collector.extractFeatures();

    ImageMatcher matcher(collector, debug);
    matcher.align_all();

    //    ImageComposer blender(mosaic_name, file_type, cache_size, tile_size, draft, qtree, verbose, grayscale);
    //    blender.blend();

}
